"""
This module defines common interface PaddleLite.
"""
import os
import numpy as np

class PaddlePaddlePredictor:
    """
    PaddlePaddle推理接口包装类，提供基于PaddlePaddle静态图的模型推理功能。
    封装了模型加载、输入设置、推理执行和输出获取的完整流程。
    """

    def __init__(self):
        import paddle as pd
        import paddle.fluid as fluid
        from paddle.fluid import debugger
        from paddle.fluid import core
        pd.enable_static()
        self.place = fluid.CPUPlace()
        self.exe = fluid.Executor(self.place)

    def load(self, model_dir):
        """
        从指定路径加载PaddlePaddle模型。
        
        Args:
            model_dir (str): 模型文件所在目录路径。
            
        Returns:
            None
            
        Notes:
            支持加载组合模型（含model和params文件）或非组合模型，
            自动识别模型文件格式并加载推理程序、输入输出信息。
        """
        import paddle as pd
        import paddle.fluid as fluid
        pd.enable_static()
        # model_dir = j["model"]
        # print(colorize('Loading model: {}'.format(model_dir), fg='green'))
        program = None
        feed = None
        fetch = None
        if os.path.exists(model_dir + "/params"):
            [program, feed, fetch] = fluid.io.load_inference_model(
                model_dir, self.exe, model_filename='model', params_filename="params")
        else:
            print("not combined")
            [program, feed, fetch] = fluid.io.load_inference_model(model_dir, self.exe)

        self.program = program
        self.feed = feed
        self.fetch = fetch
        # print(self.program)
        self.inputs = [None] * len(self.feed)

    def set_input(self, data, index):
        """
        设置模型输入数据。
        
        Args:
            data (numpy.ndarray): 输入数据数组。
            index (int): 输入数据对应的输入端口索引。
            
        Returns:
            None
        """
        self.inputs[index] = data

    def run(self):
        """
        执行模型推理计算。
        
        Returns:
            list: 推理结果列表，每个元素为PaddlePaddle的Tensor对象。
            
        Notes:
            将输入数据构建为feed字典，通过执行器运行推理程序，
            并将结果转换为numpy数组存储。
        """

        feeds = {}
        for index, _ in enumerate(self.inputs):
            feeds[self.feed[index]] = self.inputs[index]

        self.results = self.exe.run(program=self.program,
                                    feed=feeds, fetch_list=self.fetch, return_numpy=False)
        self.outputs = []
        for res in self.results:
            self.outputs.append(np.array(res))
        return self.results

    def get_output(self, index):
        """
        获取指定索引的推理输出结果。
        
        Args:
            index (int): 输出结果的索引。
            
        Returns:
            numpy.ndarray: 对应索引的输出结果数组。
        """
        return self.results[index]


class PaddleLitePredictor:
    """
    PaddleLite推理接口包装类，提供轻量级PaddleLite模型推理功能。
    适用于边缘设备和嵌入式系统的模型部署，封装了PaddleLite的底层接口。
    """
    def __init__(self):
        self.predictor = None

    def load(self, model_dir):
        """
        从指定路径加载PaddleLite模型。
        
        Args:
            model_dir (str): 模型文件或目录路径。
            
        Returns:
            None
            
        Notes:
            支持加载组合模型（model和params文件分离）或非组合模型，
            配置支持的硬件平台（FPGA/CPU/ARM）和精度类型。
        """
        from paddlelite import Place
        from paddlelite import CxxConfig
        from paddlelite import CreatePaddlePredictor
        from paddlelite import TargetType
        from paddlelite import PrecisionType
        from paddlelite import DataLayoutType
        valid_places = (
            Place(TargetType.kFPGA, PrecisionType.kFP16, DataLayoutType.kNHWC),
            Place(TargetType.kHost, PrecisionType.kFloat),
            Place(TargetType.kARM, PrecisionType.kFloat),
        )
        config = CxxConfig()
        if os.path.exists(model_dir + "/params"):
            config.set_model_file(model_dir + "/model")
            config.set_param_file(model_dir + "/params")
        else:
            config.set_model_dir(model_dir)
        config.set_valid_places(valid_places)
        self.predictor = CreatePaddlePredictor(config)

    def set_input(self, data, index):
        input = self.predictor.get_input(index)
        input.resize(data.shape)
        input.set_data(data)

    def run(self):
        self.predictor.run()

    def get_output(self, index):
        return self.predictor.get_output(index)

